#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多模态系统测试脚本
测试问题分类、提示词模板和多模态增强功能
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from question_classifier import QuestionClassifier
from prompt_templates import prompt_templates
from multimodal_enhancer import multimodal_enhancer
from answer_generator import AnswerGenerator
import json

def test_question_classifier():
    """测试问题分类器"""
    print("=" * 50)
    print("测试问题分类器")
    print("=" * 50)
    
    classifier = QuestionClassifier()
    
    test_questions = [
        "Rust的生命周期是什么？",
        "如何使用迭代器？",
        "这个编译错误怎么解决？",
        "Rust和C++有什么区别？",
        "Rust适合什么项目？"
    ]
    
    for question in test_questions:
        try:
            result = classifier.classify(question)
            print(f"问题: {question}")
            print(f"分类: {result['label']} (置信度: {result['confidence']:.2f})")
            print("-" * 30)
        except Exception as e:
            print(f"分类失败: {e}")

def test_prompt_templates():
    """测试提示词模板系统"""
    print("\n" + "=" * 50)
    print("测试提示词模板系统")
    print("=" * 50)
    
    test_cases = [
        ("definition", "Rust的生命周期是什么？"),
        ("usage", "如何使用迭代器？"),
        ("error_debug", "这个编译错误怎么解决？"),
        ("comparison", "Rust和C++有什么区别？"),
        ("faq", "Rust适合什么项目？")
    ]
    
    for question_type, question in test_cases:
        print(f"\n问题类型: {question_type}")
        print(f"问题: {question}")
        
        # 测试基本模板
        template = prompt_templates.get_template(question_type)
        print(f"系统提示词长度: {len(template['system'])} 字符")
        print(f"用户提示词长度: {len(template['user'])} 字符")
        
        # 测试格式化模板
        formatted = prompt_templates.format_prompt(
            question_type, 
            question, 
            "相关知识点: Rust生命周期管理", 
            "对话历史: 之前讨论过所有权概念"
        )
        print(f"格式化后总长度: {len(formatted['system']) + len(formatted['user'])} 字符")
        
        # 测试多模态提示词
        multimodal = prompt_templates.get_multimodal_prompt(
            question_type, 
            question, 
            "相关知识点: Rust生命周期管理", 
            "对话历史: 之前讨论过所有权概念"
        )
        print(f"多模态提示词长度: {len(multimodal)} 字符")

def test_multimodal_enhancer():
    """测试多模态增强器"""
    print("\n" + "=" * 50)
    print("测试多模态增强器")
    print("=" * 50)
    
    test_content = """
# Rust生命周期详解

## 概念定义
生命周期是Rust中确保引用有效性的机制。

## 核心要点
- 编译时检查
- 防止悬垂引用
- 内存安全保证

## 示例代码
```rust
fn longest<'a>(x: &'a str, y: &'a str) -> &'a str {
    if x.len() > y.len() {
        x
    } else {
        y
    }
}
```

## 生命周期关系图
```mermaid
graph TD
    A[变量声明] --> B[生命周期开始]
    B --> C[使用变量]
    C --> D[生命周期结束]
    D --> E[内存释放]
```

## 对比表格
| 特性 | Rust | C++ |
|------|------|-----|
| 内存管理 | 编译时检查 | 运行时检查 |
| 安全性 | 高 | 中等 |
| 性能 | 零开销 | 有开销 |
"""
    
    enhanced = multimodal_enhancer.enhance_content(test_content)
    
    print(f"原始内容长度: {len(test_content)} 字符")
    print(f"增强后HTML长度: {len(enhanced['enhanced_html'])} 字符")
    print(f"提取的代码块数量: {len(enhanced['code_blocks'])}")
    print(f"提取的Mermaid图表数量: {len(enhanced['mermaid_diagrams'])}")
    print(f"提取的表格数量: {len(enhanced['tables'])}")
    
    # 显示提取的代码块
    for i, code_block in enumerate(enhanced['code_blocks']):
        print(f"\n代码块 {i+1}:")
        print(f"  语言: {code_block['language']}")
        print(f"  ID: {code_block['id']}")
        print(f"  代码长度: {len(code_block['code'])} 字符")
    
    # 显示提取的Mermaid图表
    for i, diagram in enumerate(enhanced['mermaid_diagrams']):
        print(f"\nMermaid图表 {i+1}:")
        print(f"  ID: {diagram['id']}")
        print(f"  类型: {diagram['type']}")
        print(f"  代码长度: {len(diagram['code'])} 字符")

def test_answer_generator():
    """测试答案生成器"""
    print("\n" + "=" * 50)
    print("测试答案生成器")
    print("=" * 50)
    
    generator = AnswerGenerator()
    
    # 模拟数据
    question = "Rust的生命周期是什么？"
    classification = {"label": "definition", "confidence": 0.9}
    retrieved_docs = [
        {
            "topic": "生命周期",
            "content": "生命周期是Rust中确保引用有效性的机制",
            "code": "fn example<'a>(x: &'a str) -> &'a str { x }"
        }
    ]
    context = [
        {"question": "什么是所有权？", "answer": "所有权是Rust的核心概念..."}
    ]
    
    try:
        # 注意：这里只是测试提示词生成，不实际调用API
        print(f"问题: {question}")
        print(f"分类: {classification['label']}")
        print(f"置信度: {classification['confidence']}")
        print(f"检索文档数量: {len(retrieved_docs)}")
        print(f"上下文长度: {len(context)}")
        
        # 测试上下文构建
        context_info = generator._build_context_info(context)
        knowledge_info = generator._build_knowledge_info(retrieved_docs)
        
        print(f"\n上下文信息长度: {len(context_info)} 字符")
        print(f"知识信息长度: {len(knowledge_info)} 字符")
        
    except Exception as e:
        print(f"测试失败: {e}")

def test_integration():
    """测试系统集成"""
    print("\n" + "=" * 50)
    print("测试系统集成")
    print("=" * 50)
    
    # 模拟完整的问答流程
    question = "Rust的生命周期是什么？"
    
    print(f"用户问题: {question}")
    
    # 1. 问题分类
    classifier = QuestionClassifier()
    try:
        classification = classifier.classify(question)
        print(f"问题分类: {classification['label']} (置信度: {classification['confidence']:.2f})")
    except Exception as e:
        print(f"分类失败: {e}")
        classification = {"label": "definition", "confidence": 0.8}
    
    # 2. 生成提示词
    try:
        multimodal_prompt = prompt_templates.get_multimodal_prompt(
            classification['label'],
            question,
            "相关知识点: Rust生命周期管理",
            "对话历史: 之前讨论过所有权概念"
        )
        print(f"生成的多模态提示词长度: {len(multimodal_prompt)} 字符")
    except Exception as e:
        print(f"提示词生成失败: {e}")
    
    # 3. 模拟答案生成（不实际调用API）
    mock_answer = """
# Rust生命周期详解

## 概念定义
生命周期是Rust中确保引用有效性的机制，用于防止悬垂引用和内存安全问题。

## 核心要点
- **编译时检查**: 在编译阶段验证引用的有效性
- **防止悬垂引用**: 确保引用指向的内存仍然有效
- **内存安全保证**: 提供内存安全而不需要垃圾回收

## 示例代码
```rust
fn longest<'a>(x: &'a str, y: &'a str) -> &'a str {
    if x.len() > y.len() {
        x
    } else {
        y
    }
}

fn main() {
    let string1 = String::from("long string is long");
    let string2 = String::from("xyz");
    let result = longest(&string1, &string2);
    println!("The longest string is '{}'", result);
}
```

## 生命周期关系图
```mermaid
graph TD
    A[变量声明] --> B[生命周期开始]
    B --> C[使用变量]
    C --> D[生命周期结束]
    D --> E[内存释放]
    
    F[引用创建] --> G[生命周期检查]
    G --> H{引用有效?}
    H -->|是| I[允许使用]
    H -->|否| J[编译错误]
```

## 相关概念
- **所有权系统**: 内存管理的基础
- **借用检查器**: 生命周期检查的核心组件
- **生命周期标注**: 显式指定生命周期参数
"""
    
    # 4. 多模态增强
    try:
        enhanced = multimodal_enhancer.enhance_content(mock_answer)
        print(f"多模态增强完成:")
        print(f"  - 代码块数量: {len(enhanced['code_blocks'])}")
        print(f"  - Mermaid图表数量: {len(enhanced['mermaid_diagrams'])}")
        print(f"  - 表格数量: {len(enhanced['tables'])}")
        print(f"  - 增强HTML长度: {len(enhanced['enhanced_html'])} 字符")
    except Exception as e:
        print(f"多模态增强失败: {e}")

def main():
    """主测试函数"""
    print("Rust知识解释智能体 - 多模态系统测试")
    print("=" * 60)
    
    try:
        test_question_classifier()
        test_prompt_templates()
        test_multimodal_enhancer()
        test_answer_generator()
        test_integration()
        
        print("\n" + "=" * 60)
        print("所有测试完成！")
        print("=" * 60)
        
    except Exception as e:
        print(f"\n测试过程中出现错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 